{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/trajectories.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from IPython.display import clear_output\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    Cylinder,\n",
    "    DiagramBuilder,\n",
    "    InverseKinematics,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    RigidTransform,\n",
    "    Role,\n",
    "    RollPitchYaw,\n",
    "    RotationMatrix,\n",
    "    Solve,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import MeshcatPoseSliders\n",
    "from manipulation.scenarios import AddIiwa, AddShape, AddWsg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interactive inverse kinematics\n",
    "\n",
    "This first cell gives us an interface that is very similar to the differential IK teleop interface that we used before.  See if you can spot any differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def teleop_inverse_kinematics():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddIiwa(plant)\n",
    "    wsg = AddWsg(plant, iiwa, welded=True)\n",
    "    plant.Finalize()\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(delete_prefix_initialization_event=False),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "\n",
    "    def my_callback(context, pose):\n",
    "        ik = InverseKinematics(plant, plant_context)\n",
    "        ik.AddPositionConstraint(\n",
    "            gripper_frame,\n",
    "            [0, 0, 0],\n",
    "            plant.world_frame(),\n",
    "            pose.translation(),\n",
    "            pose.translation(),\n",
    "        )\n",
    "        ik.AddOrientationConstraint(\n",
    "            gripper_frame,\n",
    "            RotationMatrix(),\n",
    "            plant.world_frame(),\n",
    "            pose.rotation(),\n",
    "            0.0,\n",
    "        )\n",
    "        prog = ik.get_mutable_prog()\n",
    "        q = ik.q()\n",
    "        prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "        prog.SetInitialGuess(q, q0)\n",
    "        result = Solve(ik.prog())\n",
    "        if result.is_success():\n",
    "            print(\"IK success\")\n",
    "        else:\n",
    "            print(\"IK failure\")\n",
    "        clear_output(wait=True)\n",
    "\n",
    "    meshcat.DeleteAddedControls()\n",
    "    sliders = MeshcatPoseSliders(meshcat)\n",
    "    sliders.SetPose(\n",
    "        plant.EvalBodyPoseInWorld(\n",
    "            plant_context, plant.GetBodyByName(\"body\", wsg)\n",
    "        )\n",
    "    )\n",
    "    sliders.Run(visualizer, context, my_callback)\n",
    "\n",
    "\n",
    "teleop_inverse_kinematics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This one has a collision to avoid.  Try moving it in positive $y$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def teleop_inverse_kinematics():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddIiwa(plant, \"with_box_collision\")\n",
    "    wsg = AddWsg(plant, iiwa, welded=True)\n",
    "    box = AddShape(plant, Box(0.1, 0.1, 1.0), \"box\")\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"box\", box),\n",
    "        RigidTransform([0.25, 0.0, 0.5]),\n",
    "    )\n",
    "    plant.Finalize()\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder, scene_graph, meshcat, MeshcatVisualizerParams()\n",
    "    )\n",
    "    collision = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(\n",
    "            prefix=\"collision\", role=Role.kProximity, visible_by_default=False\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "\n",
    "    def my_callback(context, pose):\n",
    "        ik = InverseKinematics(plant, plant_context)\n",
    "        ik.AddPositionConstraint(\n",
    "            gripper_frame,\n",
    "            [0, 0, 0],\n",
    "            plant.world_frame(),\n",
    "            pose.translation(),\n",
    "            pose.translation(),\n",
    "        )\n",
    "        ik.AddOrientationConstraint(\n",
    "            gripper_frame,\n",
    "            RotationMatrix(),\n",
    "            plant.world_frame(),\n",
    "            pose.rotation(),\n",
    "            0.0,\n",
    "        )\n",
    "        ik.AddMinimumDistanceLowerBoundConstraint(0.001, 0.1)\n",
    "        prog = ik.get_mutable_prog()\n",
    "        q = ik.q()\n",
    "        prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "        prog.SetInitialGuess(q, q0)\n",
    "        result = Solve(ik.prog())\n",
    "        if result.is_success():\n",
    "            print(\"IK success\")\n",
    "        else:\n",
    "            print(\"IK failure\")\n",
    "        clear_output(wait=True)\n",
    "\n",
    "    sliders = MeshcatPoseSliders(meshcat)\n",
    "    sliders.SetPose(\n",
    "        plant.EvalBodyPoseInWorld(\n",
    "            plant_context, plant.GetBodyByName(\"body\", wsg)\n",
    "        )\n",
    "    )\n",
    "    # set the initial z lower, to make the interaction interesting.\n",
    "    sliders.SetXyz([0.4, -0.2, 0.35])\n",
    "    sliders.Run(diagram, context, my_callback)\n",
    "\n",
    "\n",
    "teleop_inverse_kinematics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This one has the hand tracking a cylinder, but is allowed to touch anywhere along the cylinder.  The sliders are controlling the pose of the cylinder. Or you can set `grasp_cylinder` to `False` and just chase the robot around with a stick."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def teleop_inverse_kinematics(grasp_cylinder=True):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddIiwa(plant, \"with_box_collision\")\n",
    "    wsg = AddWsg(plant, iiwa, welded=True)\n",
    "    cylinder = AddShape(plant, Cylinder(0.02, 1.0), \"cylinder\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder, scene_graph, meshcat, MeshcatVisualizerParams()\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "    cylinder_body = plant.GetBodyByName(\"cylinder\", cylinder)\n",
    "    cylinder_frame = plant.GetFrameByName(\"cylinder\", cylinder)\n",
    "\n",
    "    def my_callback(context, pose):\n",
    "        ik = InverseKinematics(plant, plant_context)\n",
    "        ik.AddPositionConstraint(\n",
    "            cylinder_frame,\n",
    "            [0, 0, 0],\n",
    "            plant.world_frame(),\n",
    "            pose.translation(),\n",
    "            pose.translation(),\n",
    "        )\n",
    "        ik.AddOrientationConstraint(\n",
    "            cylinder_frame,\n",
    "            RotationMatrix(),\n",
    "            plant.world_frame(),\n",
    "            pose.rotation(),\n",
    "            0.0,\n",
    "        )\n",
    "        if grasp_cylinder:\n",
    "            ik.AddPositionConstraint(\n",
    "                frameB=gripper_frame,\n",
    "                p_BQ=[0, 0.1, -0.02],\n",
    "                frameA=cylinder_frame,\n",
    "                p_AQ_lower=[0, 0, -0.5],\n",
    "                p_AQ_upper=[0, 0, 0.5],\n",
    "            )\n",
    "            ik.AddPositionConstraint(\n",
    "                frameB=gripper_frame,\n",
    "                p_BQ=[0, 0.1, 0.02],\n",
    "                frameA=cylinder_frame,\n",
    "                p_AQ_lower=[0, 0, -0.5],\n",
    "                p_AQ_upper=[0, 0, 0.5],\n",
    "            )\n",
    "        else:\n",
    "            ik.AddMinimumDistanceLowerBoundConstraint(0.001, 0.1)\n",
    "        prog = ik.get_mutable_prog()\n",
    "        q = ik.q()\n",
    "        prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "        prog.SetInitialGuess(q, q0)\n",
    "        result = Solve(ik.prog())\n",
    "        if result.is_success():\n",
    "            print(\"IK success\")\n",
    "        else:\n",
    "            print(\"IK failure\")\n",
    "        clear_output(wait=True)\n",
    "\n",
    "    sliders = MeshcatPoseSliders(meshcat)\n",
    "    sliders.SetPose(\n",
    "        RigidTransform(RollPitchYaw(np.pi / 2.0, 0, 0), [0.5, 0, 0.5])\n",
    "    )\n",
    "    sliders.Run(visualizer, context, my_callback)\n",
    "\n",
    "\n",
    "# Set grasp_cylinder=False if you just want to antagonize the robot with a stick.\n",
    "teleop_inverse_kinematics(grasp_cylinder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
